# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import argparse
import datetime
import json
import random
import time
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, DistributedSampler

import datasets
import util.misc as utils
from datasets import build_dataset, get_coco_api_from_dataset
from engine import evaluate, train_one_epoch
from models import build_model
from util.misc import is_main_process

from torchinfo import summary
# from torchsummary import summary


def get_args_parser():  # 参数设置
    parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
    parser.add_argument('--lr', default=1e-4, type=float)
    parser.add_argument('--lr_backbone', default=1e-5, type=float)  # default=1e-5，如果设置冻结骨干网络，此值设置为0
    parser.add_argument('--batch_size', default=4, type=int)  # Backbone不冻结(batch_size=2适合8Gb显存，batch_size=4适合12Gb显存, mamba设置batch_size=6);BB冻结数据集小(mamba设bs=14)
    parser.add_argument('--weight_decay', default=1e-4, type=float)
    parser.add_argument('--epochs', default=500, type=int)
    parser.add_argument('--lr_drop', default=200, type=int)
    parser.add_argument('--clip_max_norm', default=0.1, type=float,
                        help='gradient clipping max norm')
    parser.add_argument('--base_bs', default=4, type=int)  # 设置基础的batch_size，避免不同 batch_size下由于 lr尺度不同导致精度下降

    # Model parameters
    parser.add_argument('--frozen_weights', type=str, default=None,
                        help="Path to the pretrained model. If set, only the mask head will be trained")
    # * Backbone
    parser.add_argument('--backbone', default='resnet50', type=str,
                        help="Name of the convolutional backbone to use")
    parser.add_argument('--dilation', action='store_true',  # 是否在最后的卷积块（DC5）中使用空洞卷积，它用于增加卷积层的感受野
                        help="If true, we replace stride with dilation in the last convolutional block (DC5)")
    parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),  # 位置编码的模式，分为不学习的余弦编码和可学习的编码
                        help="Type of positional embedding to use on top of the image features")
    parser.add_argument('--no_sync_bn', dest='sync_bn', action='store_false')  # 是否在Backbone多卡训练时同步Batch Norm，默认同步

    # * Transformer
    parser.add_argument('--enc_layers', default=6, type=int,
                        help="Number of encoding layers in the transformer")
    parser.add_argument('--dec_layers', default=6, type=int,
                        help="Number of decoding layers in the transformer")
    parser.add_argument('--dim_feedforward', default=2048, type=int,  # Transformer中前馈网络层（FFN）输入通道数
                        help="Intermediate size of the feedforward layers in the transformer blocks")
    parser.add_argument('--hidden_dim', default=256, type=int,  # Transformer模型中输入的维度，即输入张量的channel数
                        help="Size of the embeddings (dimension of the transformer)")
    parser.add_argument('--dropout', default=0.1, type=float,
                        help="Dropout applied in the transformer")
    parser.add_argument('--nheads', default=8, type=int,  # 每个多头自注意力机制中的注意力头数
                        help="Number of attention heads inside the transformer's attentions")
    parser.add_argument('--num_queries', default=30, type=int,  # decoder 的 queries
                        help="Number of query slots")
    parser.add_argument('--pre_norm', action='store_true')  # 是否在Transformer块中使用“pre_norm”结构
    parser.add_argument('--mamba_in', action='store_true')  # 是否在Transformer块中使用mamba

    # * Segmentation
    parser.add_argument('--masks', action='store_true',
                        help="Train segmentation head if the flag is provided")

    # Loss
    parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false',  # 辅助解码损失，指在模型中的每个解码层级（或称为解码器的每个层级）都计算损失，可以帮助模型更好地学习特征的表示
                        help="Disables auxiliary decoding losses (loss at each layer)")
    # * Matcher
    parser.add_argument('--set_cost_class', default=1, type=float,
                        help="Class coefficient in the matching cost")
    parser.add_argument('--set_cost_bbox', default=5, type=float,
                        help="L1 box coefficient in the matching cost")
    parser.add_argument('--set_cost_giou', default=2, type=float,
                        help="giou box coefficient in the matching cost")
    # * Loss coefficients  # 定义损失函数中各项损失的权重,调节其在总体损失中的程贡献度
    parser.add_argument('--mask_loss_coef', default=1, type=float)  # 掩码损失的权重
    parser.add_argument('--dice_loss_coef', default=1, type=float)  # Dice 损失的权重，是一种常用于图像分割任务的损失函数，用于衡量模型输出的分割结果与真实分割间的之相似度，
    parser.add_argument('--bbox_loss_coef', default=5, type=float)  # 边界框损失的权重
    parser.add_argument('--giou_loss_coef', default=2, type=float)  # GIoU 损失的权重
    parser.add_argument('--eos_coef', default=0.1, type=float,  # 相对分类损失中无对象类别的权重
                        help="Relative classification weight of the no-object class")

    # dataset parameters
    parser.add_argument('--dataset_file', default='coco')
    parser.add_argument('--coco_path', type=str)
    parser.add_argument('--coco_panoptic_path', type=str)
    parser.add_argument('--remove_difficult', action='store_true')

    parser.add_argument('--output_dir', default='',
                        help='path where to save, empty for no saving')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=42, type=int)
    parser.add_argument('--resume', default='', help='resume from checkpoint')  # 是否从已有checkpoint开始训练
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--eval', action='store_true')  # action='store_true' 的意思是，如果在运行脚本中提供了 --eval，例如“python XX.py --eval”，那么 args.eval 就会被设置为 True，否则为 False
    parser.add_argument('--num_workers', default=2, type=int)

    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,  # default=1
                        help='number of distributed processes')
    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
    return parser


def main(args):
    utils.init_distributed_mode(args)
    print("git:\n  {}\n".format(utils.get_sha()))

    if args.frozen_weights is not None:
        assert args.masks, "Frozen training is meant for segmentation only"  # 如果 args.masks 不为 True，则会触发断言错误-
    print(args)                                                              # -意思是冻结训练仅用于分割任务：训练分两阶段①训练主干②训练分割头

    device = torch.device(args.device)  # 设置运行设备 CPU/GPU

    # fix the seed for reproducibility
    seed = args.seed + utils.get_rank()  # 设计随机种子 以保证每次重新训练模型需要初始化模型参数的时候能够得到相同的初始化参数，从而达到稳定复现训练结果的目的
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
# ----------------重新阅读到此------重新debug到此--------------先阅读build_model构建模型部分吧，参考CSDN博客
    model, criterion, postprocessors = build_model(args)  # 构建模型、损失函数、后处理方法
    model.to(device)

    model_without_ddp = model  # 另存模型
    if args.distributed:
        if args.sync_bn and args.world_size > 1 and args.lr_backbone > 0:  # 添加 SyncBatchNorm 支持，设置多卡同步骨干网络的 BatchNorm（如果解冻骨干）
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
            print("Sync BatchNorm has been set.")
        elif args.sync_bn:
            print("Sync_bn is not supported on a single GPU or when not distributed or frozen backbone.")
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])  # 设置多GPU的DDP模式
        model_without_ddp = model.module

    args.lr = np.sqrt((args.batch_size/args.base_bs)*args.world_size) * args.lr  # 设置平方根增加学习率，目的保持梯度的方差  # 学习率设置为根据并行GPU的数变换，因为 batch_size 的增大会导致 update 次数的减少，因此需要增大lr
    args.lr_backbone = np.sqrt((args.batch_size/args.base_bs)*args.world_size) * args.lr_backbone  # 骨干网络的 lr 也同理
    print(f"New learning rate: {args.lr:.10f}")

    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)  # 计算有梯度（可学习）的参数数量
    print('number of params:', n_parameters)
    summary(model, input_size=(args.batch_size, 3, 500, 375), device="cuda" if torch.cuda.is_available() else "cpu")  # 打印模型结构

    param_dicts = [  # 设置参数字典，区分backbone网络和其余网络部分的参数，方便读取
        {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
        {
            "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
            "lr": args.lr_backbone,
        },
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr,  # 设置优化器，weight_decay用于控制模型参数更新过程中的正则化项，避免模型过拟合
                                  weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)  # 采用StepLR设置学习率调度，在每经过lr_drop个epoch后，将学习率乘衰减因子（默认0.1）进行衰减

    dataset_train = build_dataset(image_set='train', args=args)  # 读取训练集和数据集对象
    dataset_val = build_dataset(image_set='val', args=args)

    if args.distributed:  # 根据是否分布式训练，设置不同采样器（只保存样本的索引，它们并不负责存储或管理样本数据本身）
        sampler_train = DistributedSampler(dataset_train)  # 对训练集使用 D..S.. 进行采样，这样每个进程（GPU）将会从自己负责的数据片段中采样
        sampler_val = DistributedSampler(dataset_val, shuffle=False)  # 同样使用D..S.. 进行采样，shuffle=False表示不对验证集进行随机打乱
    else:
        sampler_train = torch.utils.data.RandomSampler(dataset_train)  # 对训练集进行随机采样，每个 epoch都会对数据进行随机打乱
        sampler_val = torch.utils.data.SequentialSampler(dataset_val)  # 对验证集进行顺序采样，按顺序逐个读取验证集中的样本

    batch_sampler_train = torch.utils.data.BatchSampler(  # 创建批量采样器
        sampler_train, args.batch_size, drop_last=True)  # 设置 drop_last=True 表示丢弃最后一个不足批次大小的样本

    data_loader_train = DataLoader(dataset_train, batch_sampler=batch_sampler_train,  # 创建数据加载器
                                   collate_fn=utils.collate_fn, num_workers=args.num_workers)  # collate_fn对每个批次进行数据处理，传入nested_tensor_from_tensor_list()在此转换为NestedTensor(tensor, mask)
    data_loader_val = DataLoader(dataset_val, args.batch_size, sampler=sampler_val,
                                 drop_last=False, collate_fn=utils.collate_fn, num_workers=args.num_workers)  # num_workers指定数据加载器的工作进程数量

    if args.dataset_file == "coco_panoptic":  # 根据不同的数据集配置构建相应的 COCO 数据集对象，并将其转换为 COCO API 对象
        # We also evaluate AP during panoptic training, on original coco DS
        coco_val = datasets.coco.build("val", args)
        base_ds = get_coco_api_from_dataset(coco_val)
    else:
        base_ds = get_coco_api_from_dataset(dataset_val)  # base_ds 是 COCO 数据集的基础对象，用于提供数据集的基本信息，比如图像路径、标注信息等。它通常用于评估函数中，作为评估过程的基础数据源
                                                          # 简而言之，base_ds 主要关注的是数据集的管理和评估，而 data_loader_val 主要关注的是数据的批量加载和预处理，二者在功能和用途上有所不同
    if args.frozen_weights is not None:  # 如果 frozen_weights 被赋值，加载预训练的模型参数（冻结训练仅用于分割任务）
        checkpoint = torch.load(args.frozen_weights, map_location='cpu')  # 加载指定路径下的预训练权重文件（映射到CPU）
        model_without_ddp.detr.load_state_dict(checkpoint['model'])  # 将权重参数（字典格式）赋值给模型 DETR 部分

    output_dir = Path(args.output_dir)
    if args.resume:  # 是否从已有checkpoint开始继续训练
        if args.resume.startswith('https'):  # 如果参数是以 "https" 开头的 URL
            checkpoint = torch.hub.load_state_dict_from_url(
                args.resume, map_location='cpu', check_hash=True)
        else:  # 如果参数是本地文件路径
            checkpoint = torch.load(args.resume, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])  # 从checkpoint加载模型权重参数
        if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:  # 如果不是在评估模式下
            optimizer.load_state_dict(checkpoint['optimizer'])  # 且加载的模型包含了优化器、学习率调度器、轮数，则加载
            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
            args.start_epoch = checkpoint['epoch'] + 1

    if args.eval:  # 如果评估，评估模型
        test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
                                              data_loader_val, base_ds, device, args.output_dir)  # base_ds 是数据集对象本身，而 data_loader_val 是一个负责从数据集中加载数据的工具
        if args.output_dir:
            utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
        return

    best_mAP = 0.0
    last_checkpoint_path = None
    best_checkpoint_path = None
    print("is_main_process:", is_main_process())
    print("is_train_backbone:", args.lr_backbone > 0)
    print("is_mamba_in:", args.mamba_in)
    print("Start training")
    start_time = time.time()
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            sampler_train.set_epoch(epoch)
        train_stats = train_one_epoch(  # 进行一个周期的训练，train_stats 用于存储这个周期的训练统计信息
            model, criterion, data_loader_train, optimizer, device, epoch,
            args.clip_max_norm)
        lr_scheduler.step()  # 学习率调度器会根据预设的策略调整学习率

        if args.output_dir and utils.is_main_process():  # 保存checkpoint文件  # 不用担心遗弃其余进程权重，因为经过train_one_epoch后，已同步所有 GPU 模型的权重是相同的
            # Define paths for checkpoints
            new_last_checkpoint_path = output_dir / f"checkpoint_last_epoch{epoch+1}.pth"
            checkpoint_paths = [new_last_checkpoint_path]
            # extra checkpoint before LR drop and every 100 epochs
            if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
                checkpoint_paths.append(output_dir / f'checkpoint_epoch{epoch+1:04}.pth')
            # Save checkpoints
            for checkpoint_path in checkpoint_paths:
                utils.save_on_master({
                    'model': model_without_ddp.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'lr_scheduler': lr_scheduler.state_dict(),
                    'epoch': epoch,
                    'args': args,
                }, checkpoint_path)
            # Delete previous last checkpoint if exists
            # if isinstance(last_checkpoint_path, Path) and last_checkpoint_path.exists():
            if last_checkpoint_path is not None and last_checkpoint_path.exists():
                last_checkpoint_path.unlink()
            # Update last_checkpoint_path to the new last checkpoint
            last_checkpoint_path = new_last_checkpoint_path

        test_stats, coco_evaluator = evaluate(  # 评估模型
            model, criterion, postprocessors, data_loader_val, base_ds, device, args.output_dir
        )

        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},  # 收集日志统计信息
                     **{f'test_{k}': v for k, v in test_stats.items()},
                     'epoch': epoch,
                     'n_parameters': n_parameters}

        if args.output_dir and utils.is_main_process():  # 记录日志信息
            with (output_dir / "log.txt").open("a") as f:  # 打开（或创建） log.txt 文件，进入追加模式
                f.write(json.dumps(log_stats) + "\n")  # 使用 json.dumps(log_stats) 将 log_stats 字典转换为 JSON 字符串，并写入文件，后跟换行符

            # for evaluation logs
            if coco_evaluator is not None:
                (output_dir / 'eval').mkdir(exist_ok=True)
                if "bbox" in coco_evaluator.coco_eval:
                    filenames = ['latest.pth']
                    if epoch % 50 == 0:  # 每 50 个周期保存一次额外的文件
                        filenames.append(f'{epoch:03}.pth')
                    for name in filenames:
                        torch.save(coco_evaluator.coco_eval["bbox"].eval,
                                   output_dir / "eval" / name)

                    current_mAP = coco_evaluator.coco_eval['bbox'].stats[0]
                    if current_mAP >= best_mAP:  # 检查并保存新的最优checkpoint
                        new_best_checkpoint_path = output_dir / f"checkpoint_best_epoch{epoch+1}.pth"
                        utils.save_on_master({
                            'model': model_without_ddp.state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'lr_scheduler': lr_scheduler.state_dict(),
                            'epoch': epoch,
                            'args': args,
                        }, new_best_checkpoint_path)
                        if best_checkpoint_path is not None and best_checkpoint_path.exists():  # 删除之前的最优checkpoint
                            best_checkpoint_path.unlink()
                        best_mAP = current_mAP
                        best_checkpoint_path = new_best_checkpoint_path



    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))


if __name__ == '__main__':
    parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)
